import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from iqual import iqualnlp, evaluation, crossval
import matplotlib.pyplot as plt
import seaborn as sns
from convokit import Corpus, download

# Download the Stanford Politeness Corpus (same as in Getting Started)
print("Downloading the Stack Exchange Politeness Corpus...")
corpus = Corpus(filename=download("stack-exchange-politeness-corpus"))
print(f"Download complete. Corpus contains {len(corpus.get_utterance_ids())} utterances.")

# Process the corpus into a DataFrame (using the same code from Getting Started)
rows = []
for utt in corpus.iter_utterances():
    request_text = utt.text
    meta = utt.meta
    if 'Binary' in meta:
        binary_label = meta['Binary']
        is_polite = 1 if binary_label == 1 else 0
        
        if len(request_text.split('.')) > 1:
            parts = request_text.split('.', 1)
            context_text = parts[0] + '.'
            request_text = parts[1].strip()
        else:
            context_text = "Stack Exchange discussion:"
        
        rows.append({
            'question': context_text,
            'answer': request_text,
            'is_polite': is_polite
        })        

# Convert to DataFrame
data = pd.DataFrame(rows)
data = data[data['question'].str.len() > 0]
data = data[data['answer'].str.len() > 0]


print(f"Created DataFrame with {len(data)} records")
print(f"Polite requests: {sum(data['is_polite'] == 1)}")
print(f"Impolite requests: {sum(data['is_polite'] == 0)}")

# Split the data into training and testing sets
train_data, test_data = train_test_split(data, test_size=0.3, random_state=42)
X_train = train_data[['question', 'answer']]
y_train = train_data['is_polite']
X_test = test_data[['question', 'answer']]
y_test = test_data['is_polite']

question_col = 'question'
answer_col   = 'answer'

# Define vectorizer options
# Option 1: Use scikit-learn vectorizers [https://scikit-learn.org/stable/api/sklearn.feature_extraction.html]

SKLEARN_QA_PARAMS = {
    "Input": {
        "question": {
            "vectorizer": {
                "model": ['TfidfVectorizer', 'CountVectorizer'],
                "env": ['scikit-learn'],
                'max_features': np.arange(1000, 6000, 1000),
                'ngram_range': [(1, 1), (1, 2)]
            }
        },
        "answer": {
            "vectorizer": {
                "model": ['TfidfVectorizer', 'CountVectorizer'],
                "env": ['scikit-learn'],
                'max_features': np.arange(1500, 6500, 1000),
                'ngram_range': [(1, 1), (1, 2)]
            }
        }
    }
}

# Option 2: Use only Answer text with scikit-learn vectorizers [https://www.sbert.net/docs/package_reference/sentence_transformer/SentenceTransformer.html]
SKLEARN_A_PARAMS = {
    "Input": {
        "question": ["drop"], 
        "answer": {
            "vectorizer": {
                "model": ['TfidfVectorizer', 'CountVectorizer'],
                "env": ['scikit-learn'],
                'max_features': np.arange(1500, 6500, 1000)
            }
        }
    }
}

# Option 3: Use sentence transformers for both question and answer
SBERT_PARAMS = {
    "Input": {
        "question": {
            "vectorizer": {
                "model": ['distiluse-base-multilingual-cased-v2'],
                "env": ['sentence-transformers']
            }
        },
        "answer": {
            "vectorizer": {
                "model": ['distiluse-base-multilingual-cased-v2'],
                "env": ['sentence-transformers']
            }
        }
    }
}

# Define classifier options [https://scikit-learn.org/stable/api/sklearn.linear_model.html]
LOGISTIC_PARAMS = {
    "Classifier": {
        "model": ["LogisticRegression"],
        "C": [0.001, 0.01, 0.1, 1, 10],
        'penalty': ['l1', 'l2'],
        'solver': ['saga']
    }
}
# [https://scikit-learn.org/stable/api/sklearn.ensemble.html]
RANDOM_FOREST_PARAMS = {
    "Classifier": {
        "model": ["RandomForestClassifier"],
        "n_estimators": [100, 200],
        'max_depth': [5, 10, 20]
    }
}

SVC_PARAMS = {
    "Classifier": {
        "model": ["SVC"],
        "C": [0.01, 0.1, 1, 10],
        'kernel': ['linear', 'rbf'],
        'probability': [True]
    }
}

# Combine vectorization and classifier parameters
VECTORIZATION_PARAMS = [SKLEARN_QA_PARAMS, SKLEARN_A_PARAMS, SBERT_PARAMS]
CLASSIFIER_PARAMS = [LOGISTIC_PARAMS, RANDOM_FOREST_PARAMS, SVC_PARAMS]

# Create all combinations of parameters
ALL_PARAM_COMBINATIONS = [{**vect_params, **clf_params} for vect_params in VECTORIZATION_PARAMS for clf_params in CLASSIFIER_PARAMS]

# Convert the nested parameter dictionary to the format expected by scikit-learn
CV_SEARCH_PARAMS = [crossval.convert_nested_params(p) for p in ALL_PARAM_COMBINATIONS]
print(f"Number of hyperparameter configurations: {crossval.count_hyperparameters(CV_SEARCH_PARAMS)}")

## Create a standard KFold splitter
from sklearn.model_selection import KFold
n_splits = 3
kfold    = KFold(n_splits=n_splits, shuffle=True, random_state=42)

print("Running cross-validation with hyperparameter search...")

cv_dict = hp_model.cross_validate_fit(
    X_train,
    y_train,
    search_parameters=CV_SEARCH_PARAMS,
    cv_method='RandomizedSearchCV',    # Options: 'GridSearchCV', 'RandomizedSearchCV'
    n_iter=10, 
    # NOTE: Using `RandomizedSearchCV` and 10 iterations for demonstration purposes.
    
    scoring=scoring_dict,        # Dictionary of scoring metrics
    refit='f1',                  # Metric to use for selecting the best model
    n_jobs=8,                   # Use all available processors
    cv_splits=kfold,                 # Number of cross-validation folds
)

# Display cross-validation results of best score
print(f"Average test score across folds: {cv_dict['avg_test_score']:.3f}")

# Get the best parameters from cross-validation
best_params_df = pd.DataFrame([hp_model.cv.get_best_params()])
# Extract and display the most important parameters
selected_params = {}
for param, value in best_params.items():
    if 'model' in param or 'max_features' in param or 'C' in param or 'n_estimators' in param:
        selected_params[param] = value

selected_params['decision_threshold'] = hp_model.get_params()['steps'][-1][-1].threshold

print("Best Model Configuration:")
for param, value in selected_params.items():
    print(f"  {param}: {value}")

# Initialize arrays to store CV predictions
cv_train_pred = np.zeros(len(X_train))
cv_train_proba = np.zeros(len(X_train))


# Get CV splits
cv_splits = list(kfold.split(X_train))

# For each fold
for fold_idx, (train_idx, val_idx) in enumerate(cv_splits):
    print(f"Processing fold {fold_idx + 1}/{n_splits}")
    
    # Get fold training and validation data
    X_fold_train = X_train.iloc[train_idx]
    y_fold_train = y_train.iloc[train_idx]
    X_fold_val = X_train.iloc[val_idx]
    
    # Create and train model with best parameters
    fold_model = iqualnlp.Model()
    fold_model.add_text_features(question_col, answer_col)
    fold_model.add_classifier()
    fold_model.add_threshold(scoring_metric='f1')  # Same as refit parameter
    fold_model.compile()
    
    # Set best parameters and fit
    fold_model.set_params(**best_params)
    fold_model.fit(X_fold_train, y_fold_train)
    
    # Predict on validation fold
    cv_train_pred[val_idx] = fold_model.predict(X_fold_val)
    cv_train_proba[val_idx] = fold_model.predict_proba(X_fold_val)

from sklearn.metrics import (
    classification_report, 
    confusion_matrix, 
    accuracy_score, 
    precision_score, 
    recall_score, 
    f1_score
)

# Calculate metrics on fold-by-fold predictions
fold_metrics = {
    'Accuracy': accuracy_score(y_train, cv_train_pred),
    'Precision': precision_score(y_train, cv_train_pred),
    'Recall': recall_score(y_train, cv_train_pred),
    'F1 Score': f1_score(y_train, cv_train_pred)
}

print("Fold-by-Fold Cross-Validation Performance Metrics:")
for metric, value in fold_metrics.items():
    print(f"  {metric}: {value:.4f}")

print("\nDetailed Classification Report:")
print(classification_report(y_train, cv_train_pred))

# Create confusion matrix
cm = confusion_matrix(y_train, cv_train_pred)
plt.figure(figsize=(8, 6))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
            xticklabels=['Impolite', 'Polite'],
            yticklabels=['Impolite', 'Polite'])
plt.xlabel('Predicted')
plt.ylabel('Actual')
plt.title('Confusion Matrix (Fold-by-Fold CV)')
plt.show()

# Train final model on all training data
final_model = iqualnlp.Model()
final_model.add_text_features(question_col, answer_col)
final_model.add_classifier()
final_model.add_threshold(scoring_metric='f1')
final_model.compile()
final_model.set_params(**best_params)
final_model.fit(X_train, y_train)

# Predict on test set
test_pred = final_model.predict(X_test)
test_proba = final_model.predict_proba(X_test)

# Calculate metrics
test_metrics = {
    'Accuracy': accuracy_score(y_test, test_pred),
    'Precision': precision_score(y_test, test_pred),
    'Recall': recall_score(y_test, test_pred),
    'F1 Score': f1_score(y_test, test_pred)
}

print("Test Set Performance Metrics:")
for metric, value in test_metrics.items():
    print(f"  {metric}: {value:.4f}")

print("\nDetailed Classification Report:")
print(classification_report(y_test, test_pred))

# Create confusion matrix for test predictions
cm_test = confusion_matrix(y_test, test_pred)
plt.figure(figsize=(8, 6))
sns.heatmap(cm_test, annot=True, fmt='d', cmap='Blues', 
            xticklabels=['Impolite', 'Polite'],
            yticklabels=['Impolite', 'Polite'])
plt.xlabel('Predicted')
plt.ylabel('Actual')
plt.title('Confusion Matrix (Test Set)')
plt.show()

# Create a DataFrame to compare CV and test metrics
comparison_df = pd.DataFrame({
    'Metric': list(fold_metrics.keys()),
    'Cross-Validation': list(fold_metrics.values()),
    'Test Set': list(test_metrics.values())
})

print(comparison_df)

# Plot comparison
plt.figure(figsize=(10, 6))
comparison_df.set_index('Metric').plot(kind='bar', rot=0)
plt.ylabel('Score')
plt.title('Cross-Validation vs Test Performance')
plt.ylim(0, 1)
plt.grid(axis='y', linestyle='--', alpha=0.7)
plt.show()

# Create a DataFrame with test examples and their predictions
results_df = pd.DataFrame({
    'Question': X_test['question'],
    'Answer': X_test['answer'],
    'True Label': y_test,
    'Predicted Label': test_pred,
    'Probability': test_proba
})

# Calculate prediction error (0 = correct, 1 = incorrect)
results_df['Error'] = (results_df['True Label'] != results_df['Predicted Label']).astype(int)

decision_threshold_final = final_model.get_params()['steps'][-1][-1].threshold

plt.figure(figsize=(10, 6))
sns.histplot(data=results_df, x='Probability', hue='True Label', bins=20, alpha=0.6)
plt.axvline(x=decision_threshold_final, color='black', linestyle='--', alpha=0.7, label='Decision Boundary')
plt.title('Probability Distribution by True Label')
plt.xlabel('Predicted Probability (Polite)')
plt.ylabel('Count')
plt.legend(title='True Label', labels=['Impolite', 'Polite'])
plt.grid(axis='y', linestyle='--', alpha=0.5)
plt.show()

print(f"Best model parameters were selected from {crossval.count_hyperparameters(CV_SEARCH_PARAMS)} hyperparameter combinations")
print(f"Best model F1 score in cross-validation: {cv_dict['avg_test_score']:.4f}")
print(f"Final model F1 score on test set: {test_metrics['F1 Score']:.4f}")

# Extract the key components of the best model
best_vectorizers = []
for param, value in best_params.items():
    if 'vectorizer__model' in param and isinstance(value, str):
        component = param.split('__')[1]  # question or answer
        best_vectorizers.append(f"{component}: {value}")

best_classifier = best_params.get('Classifier__model', 'Unknown')

print("\n" + '-'*50)
print("\nBest Model Configuration:")
print("\n" + '-'*50 + '\n')
print(f"Vectorizer(s): {', '.join(best_vectorizers)}")
print(f"Classifier: {best_classifier}")
print(f"Decision Boundary: {decision_threshold_final:.2f}")